package com.sjr.easy.chat.config;


import com.sjr.easy.chat.utils.TraceUtils;

import org.aopalliance.aop.Advice;
import org.aopalliance.intercept.MethodInterceptor;
import org.aopalliance.intercept.MethodInvocation;
import org.apache.commons.lang3.ArrayUtils;
import org.apache.commons.lang3.StringUtils;
import org.springframework.amqp.core.Message;
import org.springframework.amqp.core.MessageProperties;
import org.springframework.amqp.rabbit.config.SimpleRabbitListenerContainerFactory;
import org.springframework.amqp.rabbit.core.RabbitTemplate;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.beans.factory.annotation.Autowired;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;


public class TraceRabbitMQConfiguration implements InitializingBean {

    @Autowired
    private RabbitTemplate rabbitTemplate;

    @Autowired
    private SimpleRabbitListenerContainerFactory simpleRabbitListenerContainerFactory;

    @Override
    public void afterPropertiesSet() {

        // 消息发送 trace
        rabbitTemplate.addBeforePublishPostProcessors(message -> {
            if (message == null) {
                return null;
            }
            String traceId = TraceUtils.currentTraceId();
            if (StringUtils.isEmpty(traceId)) {
                traceId = TraceUtils.generateId();
            }
            message.getMessageProperties().setHeader(TraceUtils.RABBIT_MQ_TRACE_ID_KEY, traceId);
            return message;
        });


        // 消息消费 trace
        Advice[] adviceChain = simpleRabbitListenerContainerFactory.getAdviceChain();
        Advice[] newAdviceChain = new Advice[ArrayUtils.getLength(adviceChain) + 1];
        int i = 0;
        for (; adviceChain != null && i < adviceChain.length; i++) {
            newAdviceChain[i] = adviceChain[i];
        }
        newAdviceChain[i] = new MethodInterceptor() {
            @Nullable
            @Override
            public Object invoke(@Nonnull MethodInvocation invocation) throws Throwable {
                Object arg1 = invocation.getArguments()[1];
                if (arg1 instanceof Message msg) {
                    MessageProperties messageProperties = msg.getMessageProperties();
                    String traceId = messageProperties.getHeader(TraceUtils.RABBIT_MQ_TRACE_ID_KEY);
                    if (StringUtils.isEmpty(traceId)) {
                        traceId = TraceUtils.generateId();
                    }
                    return TraceUtils.runWithTracer(invocation::proceed, new TraceUtils.TraceInfo(traceId, TraceUtils.generateId()));
                } else {
                    return invocation.proceed();
                }
            }
        };
        simpleRabbitListenerContainerFactory.setAdviceChain(newAdviceChain);
    }
}
